from gensim.models.callbacks import CallbackAny2Vec


# 自定义回调类，用于打印训练进度和损失值
class EpochLogger(CallbackAny2Vec):
    """
    自定义回调类，用于打印 Word2Vec 训练的进度和详细信息（如损失值）。
    """

    def __init__(self,
                 model_save_path,
                 model_save_interval=10):
        self.epoch = 0
        self.previous_loss = 0  # 保存当前累计总损失值
        self.model_save_path = model_save_path
        self.model_save_interval = model_save_interval

    def on_epoch_begin(self, model):
        print(f"开始第 {self.epoch + 1} 轮训练...")

    def on_epoch_end(self, model):
        # 获取当前轮次的累计总损失值
        current_loss = model.get_latest_training_loss()

        loss_diff = current_loss - self.previous_loss
        print(f"完成第 {self.epoch + 1} 轮训练，当前损失值: {loss_diff}")

        # 更新 previous_loss
        self.previous_loss = current_loss
        self.epoch += 1

        # 保存模型, 每10轮保存一次
        if self.epoch % self.model_save_interval == 0:
            model.save(f"{self.model_save_path}/word2vec_epoch_{self.epoch}.model")

    def on_train_begin(self, model):
        print("Word2Vec 训练开始...")

    def on_train_end(self, model):
        print("Word2Vec 训练结束。")